import os
import numpy as np
import paddle
from paddle.vision.transforms import Compose, ToTensor
from utils.dataset import LaneDataset


def class_statistic(num_classes):

    train_dataset = LaneDataset("train.csv", transform=Compose([ToTensor()]))

    data_batch = paddle.io.DataLoader(train_dataset, batch_size=4, shuffle=True, drop_last=False)

    class_result = {i : 0 for i in range(num_classes)}
    for item in data_batch:
        label_image = item[1].numpy()
        for cls in range(num_classes):
            class_result[cls] += np.sum(label_image == cls)

    return class_result


def class_graph_statistic(num_classes):

    train_dataset = LaneDataset("train.csv", transform=Compose([ToTensor()]))

    data_batch = paddle.io.DataLoader(train_dataset, batch_size=1, shuffle=True, drop_last=False)

    class_in_graph_result = {i : 0 for i in range(num_classes)}
    for item in data_batch:
        label_image = item[1].numpy()
        for cls in range(num_classes):
            if np.count_nonzero(label_image == cls) != 0:
                class_in_graph_result[cls] += 1

    return class_in_graph_result


# def get_file(root_path,all_files=[]):
#     '''
#     递归函数，遍历该文档目录和子目录下的所有文件，获取其path
#     '''
#     files = os.listdir(root_path)
#     for file in files:
#         if not os.path.isdir(os.path.join(root_path, file)):   # not a dir
#             if os.path.splitext(file)[-1] == '.png':
#                 all_files.append(root_path + '/' + file)
#         else:  # is a dir
#             get_file((os.path.join(root_path, file)), all_files)
#     return all_files


if __name__ == '__main__':
    # label_files = get_file(path)
    print('class in graph', class_graph_statistic(8))
    print('class statistic', class_statistic(8))